package com.zhm.langchain4j.controller;

import com.zhm.langchain4j.AiModelClient;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;

@Slf4j
@RestController
@RequestMapping("/ai")
@CrossOrigin
public class TestController {

    @Autowired
    private AiModelClient aiModelClient;

    @GetMapping("chat")
    public String chat(@RequestParam(value = "text", required = false) String text) throws Exception {
        log.info("text:{}", text);
        return "pong " + text;
    }

    @RequestMapping("/stream")
    public SseEmitter stream(@RequestParam(value = "text") String text) throws Exception {
        SseEmitter emitter = new SseEmitter(-1L);
        emitter.onTimeout(() -> {
            log.error("onTimeout {}", text);
        });
        emitter.onCompletion(() -> {
            log.info("onCompletion {}", text);
        });
        emitter.onError((error) -> {
            log.error(error.getMessage());
        });
        log.info("text:{} connected", text);

        AIRes aiRes = new AIRes();
        aiRes.setData("");
        StreamingChatLanguageModel model = aiModelClient.getStreamModel();
        model.chat(text, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String s) {
                try {
                    aiRes.setData(aiRes.getData() + s);
                    emitter.send(aiRes);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }

            @Override
            public void onCompleteResponse(ChatResponse chatResponse) {
                try {
                    aiRes.setF(1);
                    emitter.send(aiRes);
                    emitter.complete();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }

            }

            @Override
            public void onError(Throwable throwable) {
                emitter.completeWithError(throwable);
                log.error("error", throwable);
            }
        });
        return emitter;
    }

}
